import gradio as gr
import requests


def predict_text(data):
    url = "http://localhost:8000/predict"
    # 修改请求体的键为"text"，以匹配FastAPI接口的期望
    payload = {"text": data}
    try:
        response = requests.post(url, json=payload)
        response.raise_for_status()  # 检查请求是否成功，若失败会抛出异常
        return response.json()["prediction"]
    except requests.exceptions.RequestException as e:
        # 捕获请求过程中的异常，包括网络问题、HTTP错误等
        return f"请求出错: {str(e)}"
    except KeyError as e:
        # 捕获解析响应时的KeyError异常
        return f"解析响应出错: {str(e)}"


iface = gr.Interface(
    fn=predict_text,
    inputs=gr.Textbox(label="输入数据"),
    outputs=gr.Textbox(label="预测结果"),
    title="文本分类",
    description="输入特定格式数据，通过FastAPI加载的模型进行分类"
)

iface.launch()